[프로그래머스] 모두 0으로 만들기
April 16, 2021
https://programmers.co.kr/learn/courses/30/lessons/76503#
문제
문제
각 점에 가중치가 부여된 트리가 주어집니다. 당신은 다음 연산을 통하여, 이 트리의 모든 점들의 가중치를 0으로 만들고자 합니다.
임의의 연결된 두 점을 골라서 한쪽은 1 증가시키고, 다른 한쪽은 1 감소시킵니다. 하지만, 모든 트리가 위의 행동을 통하여 모든 점들의 가중치를 0으로 만들 수 있는 것은 아닙니다. 당신은 주어진 트리에 대해서 해당 사항이 가능한지 판별하고, 만약 가능하다면 최소한의 행동을 통하여 모든 점들의 가중치를 0으로 만들고자 합니다.
트리의 각 점의 가중치를 의미하는 1차원 정수 배열 a와 트리의 간선 정보를 의미하는 edges가 매개변수로 주어집니다. 주어진 행동을 통해 트리의 모든 점들의 가중치를 0으로 만드는 것이 불가능하다면 -1을, 가능하다면 최소 몇 번만에 가능한지를 찾아 return 하도록 solution 함수를 완성해주세요. (만약 처음부터 트리의 모든 정점의 가중치가 0이라면, 0을 return 해야 합니다.)
풀이
월간 코드 챌린지에 나왔었는데, 대회 시간 내에는 풀지 못하고 이후에 해설을 보고 힌트를 얻어 풀었다. 대회때는 뭔가 위상정렬로 풀면 풀릴 것 같아서 위상정렬을 시도했는데, 아이디어는 비슷했지만 접근에서 틀린 부분이 있었다.
문제를 풀기 위해서는 위상정렬에서 하는 것 처럼 각 정점에 연결된 간선의 개수를 알아야 한다. 왜냐하면 트리의 리프노드에서부터 가중치를 계속 위로 올려서 정점 하나가 남을 때까지 진행하면 연산 횟수를 구할 수 있기 때문이다. 나는 각 정점을 양바향으로 연결하고, inDegree 값이 1인 정점을 먼저 골라 가중치를 자신의 부모노드로 올리도록 했다. 그리고 처리가 끝난 리프노드의 inDegree 값을 -1 로 만들어 트리에서 제외시키고, 방금 노드 하나를 제외시키므로써 만들어지는 새로운 트리의 리프노드를 다시한번 순회하도록 했다.
코드
{% raw %}
#include <vector>
#include <queue>
#include <cmath>
#include <iostream>
using namespace std;
long long solution(vector<int> a, vector<vector<int>> edges) {
long long answer = 0;
vector<vector<int>> adj(a.size());
vector<int> inDegree(300001, 0);
int sum = 0;
for (int i = 0; i < a.size(); i++) {
sum += a[i];
}
if (sum != 0) return -1;
for (auto edge : edges) {
adj[edge[0]].push_back(edge[1]);
adj[edge[1]].push_back(edge[0]);
inDegree[edge[1]]++;
inDegree[edge[0]]++;
}
int nodeLeft = a.size();
int size = a.size();
int i = 0;
vector<long long> accumulatedSum;
accumulatedSum.assign(a.begin(), a.end());
while (nodeLeft > 1) {
for (int i = 0; i < a.size(); i++) {
if (inDegree[i] == 1) {
inDegree[i] = -1;
for (auto next : adj[i]) {
if (inDegree[next] != -1) {
accumulatedSum[next] += accumulatedSum[i];
answer += abs(accumulatedSum[i]);
accumulatedSum[i] = 0;
inDegree[next]--;
}
}
nodeLeft--;
}
}
}
return answer;
}
{% endraw %}
